package netnotifier.builder;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.HashMap;

import netnotifier.NetNotifier;
import netnotifier.notification.IPv4Packet;
import netnotifier.notification.RawPacket;
import netnotifier.notification.TCPPacket;
import netnotifier.notification.TCPRead;
import netnotifier.notification.TCPSession;
import netnotifier.notification.TCPWrite;

public class TCPBuilder
{
	NetNotifier notifier=NetNotifier.instance;

	HashMap<Long,TCPSession> synSent=new HashMap<Long,TCPSession>();
	HashMap<Long,TCPSession> synReceived=new HashMap<Long,TCPSession>();

	// the list of all established connections, lookup by expected client seq
	HashMap<Long,TCPSession> establishedClient=new HashMap<Long,TCPSession>();
	
	// the list of all established connections, lookup by expected server seq
	HashMap<Long,TCPSession> establishedServer=new HashMap<Long,TCPSession>();
	
	
	public TCPBuilder()
	{
		try
		{
			notifier.subscribe(TCPPacket.class,this,"gotTCP");
		}
		catch(Exception e)
		{
			e.printStackTrace();
		}
	}
	
	// this method gets notified whenever a tcp packet comes along
	public void gotTCP(Object packet)
	{
		TCPPacket tcpPacket=(TCPPacket)packet;
		
		// fin=1 syn=2 rst=4 psh=8 ack=16
		
		if(tcpPacket.flags==2)gotSyn(tcpPacket);
		//else if(tcpPacket.flags==4)gotRst(tcpPacket);
		else if(tcpPacket.flags==16)gotAck(tcpPacket); // just make sure the ack flag is set
		else if(tcpPacket.flags==17)gotFinAck(tcpPacket);
		else if(tcpPacket.flags==18)gotSynAck(tcpPacket);
		//else if(tcpPacket.flags==20)gotAckRst(tcpPacket);
		else if(tcpPacket.flags==24)gotPshAck(tcpPacket);
		//else if(tcpPacket.flags==25)gotFinAckPsh(tcpPacket);
		else gotOther(tcpPacket);
		
		//System.out.println(tcpPacket.type);
	}
	
	// syn packets are the first packets in a tcp session
	// 
	// 1. create a TCPSession
	// 2. dig through the syn packet for ip information
	// 3. resolve the server hostname
	// 4. set the session state to SYN_SENT
	// 5. set the server sequence number for the session
	// 6. add the session to the synSent hash
	// 7. contribute the TCPSession to the notification system
	public void gotSyn(TCPPacket tcpPacket)
	{	
		//tcpPacket.type="syn";
		TCPSession tcpSession=new TCPSession();
		
		RawPacket rawPacket=(RawPacket)tcpPacket.components.get(0);
		tcpSession.timestamp=""+rawPacket.timestamp; // lame type conversion
		IPv4Packet ipv4Packet=null;
		for(Object o:rawPacket.contributions.toArray())
		{
			if(o instanceof IPv4Packet)
			{
				ipv4Packet=(IPv4Packet)o;
				break;
			}
		}
		if(ipv4Packet==null)
		{
			tcpSession.clientIP="none";
			tcpSession.serverIP="none";
			tcpSession.hostName="none";
		}
		else
		{
			byte[] clientIP=ipv4Packet.sourceIP;
			byte[] serverIP=ipv4Packet.destinationIP;
		
			tcpSession.clientIP=new String((clientIP[0]&0xff)+"."+(clientIP[1]&0xff)+"."+(clientIP[2]&0xff)+"."+(clientIP[3]&0xff));
			tcpSession.serverIP=new String((serverIP[0]&0xff)+"."+(serverIP[1]&0xff)+"."+(serverIP[2]&0xff)+"."+(serverIP[3]&0xff));

			try
			{
				tcpSession.hostName=InetAddress.getByName(tcpSession.serverIP).getHostName();
			}
			catch(UnknownHostException e)
			{
				tcpSession.hostName="unknown";
			}
		}
		
		tcpSession.state="SYN_SENT";
		
		// set the next expected seq from the server
		tcpSession.serverNumber=tcpPacket.sequenceNumber+1;
		synSent.put(tcpSession.serverNumber,tcpSession);
		
		tcpPacket.contributions.add(tcpSession);
		tcpSession.components.add(tcpPacket);
		
		//System.out.println("syn");
		notifier.contribute(tcpSession);
	}
	
	// the syn/ack is the first packet that the server sends in a session
	//
	// 1. look up the TCPSession by searching for the ack number in the synSent hash
	// 2. move the TCPSession from the synSent hash to the synReceived hash
	public void gotSynAck(Object packet)
	{
		TCPPacket tcpPacket=(TCPPacket)packet;
		//tcpPacket.type="syn/ack";
		
		if(synSent.containsKey(tcpPacket.acknowledgementNumber))
		{
			TCPSession session=synSent.get(tcpPacket.acknowledgementNumber);
			synSent.remove(tcpPacket.acknowledgementNumber);
			
			// set the next expected seq from the client
			session.clientNumber=tcpPacket.sequenceNumber+1;
			session.state="SYN_RECEIVED";
			session.components.add(packet);
			synReceived.put(session.clientNumber,session);
		}
		else
		{
			//System.out.println("TCPBuilder warning: unused syn/ack");
		}
	}
	
	// ack packets usually make up a majority of a TCP session
	// they are used to finish off the TCP 3 way handshake, and pass bulk data
	public void gotAck(Object packet)
	{
		TCPPacket tcpPacket=(TCPPacket)packet;
		//tcpPacket.type="ack";
		
		//System.out.print("ack .. ");
		
		// see if this is part of a synReceived session
		if(synReceived.containsKey(tcpPacket.acknowledgementNumber))
		{
			//System.out.println("part of the three way handshake");
			TCPSession session=synReceived.get(tcpPacket.acknowledgementNumber);
			synReceived.remove(tcpPacket.acknowledgementNumber);
			
			session.state="ESTABLISHED";
			session.components.add(packet);
			session.currentRead=new TCPRead();
			notifier.contribute(session.currentRead);
			
			//System.out.println("hells yes, connection established  "+session+" "+tcpPacket.dataLength+" "+Thread.currentThread());
			
			establishedServer.put(session.serverNumber,session);
			establishedClient.put(session.clientNumber,session);
		}
		// see if this is part of a write to a server (this is an ack that the server expects)
		else if(establishedClient.containsKey(tcpPacket.acknowledgementNumber))
		{
			//System.out.println("got packet from the client "+tcpPacket.sourceIP+" -> "+tcpPacket.destinationIP+"  ("+tcpPacket.acknowledgementNumber+")");
			TCPSession session=establishedClient.get(tcpPacket.acknowledgementNumber);
			
			// Increment the serverNumber
			establishedServer.remove(session.serverNumber);
			session.serverNumber=tcpPacket.sequenceNumber+tcpPacket.dataLength;
			establishedServer.put(session.serverNumber,session);
			
			//System.out.println("part of a write "+session.clientIP+" -> "+session.serverIP);
			
			// error check a bit
			if(session.currentWrite==null)return;
			session.currentWrite.append(tcpPacket.tcpData);
			session.components.add(packet);
		}
		
		// see if this is part of a read from a server
		else if(establishedServer.containsKey(tcpPacket.acknowledgementNumber))
		{
			//System.out.println("got packet from the server "+tcpPacket.sourceIP+" -> "+tcpPacket.destinationIP);
			TCPSession session=establishedServer.get(tcpPacket.acknowledgementNumber);
			
			// Increment the clientNumber
			establishedClient.remove(session.clientNumber);
			session.clientNumber=tcpPacket.sequenceNumber+tcpPacket.dataLength;
			establishedClient.put(session.clientNumber,session);
		
			//System.out.println("part of a read "+session.serverIP+" -> "+session.clientIP);
			
			if(session.currentRead==null)return;
			session.currentRead.append(tcpPacket.tcpData);
			
			// add this packet to the components for this session
			session.components.add(tcpPacket);
		}
		else
		{
			//System.out.println("TCPBuilder warning: unused ack");
		}
		
	}
	
	// this is the packet that starts a new read or write in a tcp session
	public void gotPshAck(TCPPacket tcpPacket)
	{
		//tcpPacket.type="psh/ack";
		
		//System.out.print("psh/ack ..  ");
		
		// see if this is part of a write to a server
		if(establishedClient.containsKey(tcpPacket.acknowledgementNumber))
		{
			TCPSession session=establishedClient.get(tcpPacket.acknowledgementNumber);
			
			// Increment the serverNumber
			establishedServer.remove(session.serverNumber);
			session.serverNumber=tcpPacket.sequenceNumber+tcpPacket.dataLength;
			establishedServer.put(session.serverNumber,session);
		
			if(session.currentWrite!=null)notifier.contribute(session.currentWrite);
			session.currentWrite=new TCPWrite();
			session.currentWrite.clientIP=session.clientIP;
			session.currentWrite.serverIP=session.serverIP;
			
			//System.out.println("part of a write "+session.clientIP+" -> "+session.serverIP);
			
			session.currentWrite.append(tcpPacket.tcpData);
			session.components.add(tcpPacket);
		}
		// see if this is part of a read from a server
		else if(establishedServer.containsKey(tcpPacket.acknowledgementNumber))
		{
			TCPSession session=establishedServer.get(tcpPacket.acknowledgementNumber);
			
			// Increment the clientNumber
			establishedClient.remove(session.clientNumber);
			session.clientNumber=tcpPacket.sequenceNumber+tcpPacket.dataLength;
			establishedClient.put(session.clientNumber,session);
		
			//System.out.println("part of a read "+session.serverIP+" -> "+session.clientIP);
			
			if(session.currentRead!=null)notifier.contribute(session.currentRead);
			session.currentRead=new TCPRead();
			session.currentRead.clientIP=session.clientIP;
			session.currentRead.serverIP=session.serverIP;
			session.currentRead.append(tcpPacket.tcpData);	
			session.components.add(tcpPacket);
		}
		else
		{
			//System.out.println("TCPBuilder warning: unused psh/ack");
		}
	}
	
	public void gotFinAck(TCPPacket tcpPacket)
	{
		//tcpPacket.type="fin/ack";
		
		if(establishedClient.containsKey(tcpPacket.acknowledgementNumber))
		{
			TCPSession session=establishedClient.get(tcpPacket.acknowledgementNumber);
			//System.out.println("got fin/ack packet from the client "+session.clientIP+" -> "+session.serverIP+"  ("+tcpPacket.acknowledgementNumber+")");
		
			// Increment the serverNumber
			establishedServer.remove(session.serverNumber);
			session.serverNumber=tcpPacket.sequenceNumber+tcpPacket.dataLength;
			establishedServer.put(session.serverNumber,session);
			
			// Remove this session from the set of established sessions
			establishedClient.remove(session.clientNumber);
			//establishedServer.remove(session.serverNumber);
			
			if(session.currentWrite!=null)notifier.contribute(session.currentWrite);
			if(session.currentRead!=null)notifier.contribute(session.currentRead);
			session.currentRead=null;
			session.currentWrite=null;
			session.state="CLOSED";
		}
		// see if this is part of a read from a server
		else if(establishedServer.containsKey(tcpPacket.acknowledgementNumber))
		{
			TCPSession session=establishedServer.get(tcpPacket.acknowledgementNumber);
			//System.out.println("got fin/ack packet from the server "+session.serverIP+" -> "+session.clientIP);
		
			// Increment the clientNumber
			establishedClient.remove(session.clientNumber);
			session.clientNumber=tcpPacket.sequenceNumber+tcpPacket.dataLength;
			establishedClient.put(session.clientNumber,session);
			
			// Remove this session from the set of established sessions
			//establishedClient.remove(session.clientNumber);
			establishedServer.remove(session.serverNumber);
			
			if(session.currentWrite!=null)notifier.contribute(session.currentWrite);
			if(session.currentRead!=null)notifier.contribute(session.currentRead);
			session.currentRead=null;
			session.currentWrite=null;
			session.state="CLOSED";
		}
		else
		{
			//System.out.println("TCPBuilder warning: unused fin/ack");
		}
	}
	
	public void gotOther(Object packet)
	{
		TCPPacket tcpPacket=(TCPPacket)packet;
		//System.out.println("TCPBuilder warning: unknown tcp packet  flags: "+tcpPacket.flags);
	}
}
